{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "be3d6762-59dc-4339-8118-d0f3cf158cc5",
   "metadata": {},
   "source": [
    "# 个人身份信息（PII）检测的 LLM 网关\n",
    "*作者: [Anthony Susevski](https://github.com/asusevski)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8e0d0f7-25e1-4fc7-a4c1-0c5196f776c5",
   "metadata": {},
   "source": [
    "采用大语言模型（LLM）进行企业级应用时，常见的投诉之一就是数据隐私问题，尤其是对于处理敏感数据的团队来说。尽管开源模型通常是一个不错的选择，*如果可能的话应该尝试使用*，但有时我们只想快速演示一下，或者有充分的理由使用 LLM API。在这种情况下，最好采用某种网关来处理个人身份信息（PII）数据的清洗，从而降低 PII 泄露的风险。\n",
    "\n",
    "总部位于加拿大多伦多的金融科技公司 **Wealthsimple** 已经为了这个目的 [开源了一个代码库](https://github.com/wealthsimple/llm-gateway)。在本 Notebook 中，我们将探索如何利用这个代码库，在向 LLM 提供商发出 API 调用之前，对数据进行清洗。为此，我们将使用来自[AI4Privacy](https://huggingface.co/datasets/ai4privacy/pii-masking-200k)的 PII 数据集，并使用 Cohere 的 [Command R+](https://huggingface.co/CohereForAI/c4ai-command-r-plus)模型的[免费试用 API](https://cohere.com/blog/free-developer-tier-announcement)，演示 Wealthsimple 的 PII 清洗功能。\n",
    "\n",
    "首先，请按照 [README](https://github.com/wealthsimple/llm-gateway) 中的说明进行安装：\n",
    "1. 安装 Poetry 和 Pyenv\n",
    "2. 安装 pyenv 版本 3.11.3\n",
    "3. 安装项目所需的依赖：\n",
    "```\n",
    "brew install gitleaks\n",
    "poetry install\n",
    "poetry run pre-commit install\n",
    "```\n",
    "4. 运行 `cp .envrc.example .envrc` 并用 API 密钥更新配置。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "aed2caf7-41dd-427b-b69a-18d44d554319",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from llm_gateway.providers.cohere import CohereWrapper\n",
    "from datasets import load_dataset\n",
    "import cohere\n",
    "import types\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "85d8bcd4-37a6-4add-a83d-43909bbe9c87",
   "metadata": {},
   "outputs": [],
   "source": [
    "COHERE_API_KEY = os.environ['COHERE_API_KEY']\n",
    "DATABASE_URL = os.environ['DATABASE_URL'] # default database url: \"postgresql://postgres:postgres@postgres:5432/llm_gateway\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae399eb7-f7a5-44d4-bfd5-cc38570b2360",
   "metadata": {},
   "source": [
    "## LLM 包装器\n",
    "\n",
    "包装器对象是一个简单的封装器，它在发起 API 调用之前，将“清洗器”应用到输入的提示语上。使用包装器发起请求时，我们将获得一个响应和一个 `db_record` 对象。在深入了解更多细节之前，让我们先来看一下它的实际应用。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "b0882af5-42bc-460a-a42e-9bd336dc7292",
   "metadata": {},
   "outputs": [],
   "source": [
    "wrapper = CohereWrapper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "f7f055e4-c290-4ac1-9869-2e8712355c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "example = \"Michael Smith (msmith@gmail.com, (+1) 111-111-1111) committed a mistake when he used PyTorch Trainer instead of HF Trainer.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "3f2ae809-cf5b-444c-9bff-34947575a428",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data': ['Michael Smith made a mistake by using PyTorch Trainer instead of HF Trainer.'], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 48, 'output_tokens': 14}}}\n"
     ]
    }
   ],
   "source": [
    "response, db_record = wrapper.send_cohere_request(\n",
    "    endpoint=\"generate\",\n",
    "    model=\"command-r-plus\",\n",
    "    max_tokens=25,\n",
    "    prompt=f\"{example}\\n\\nSummarize the above text in 1-2 sentences.\",\n",
    "    temperature=0.3,\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c4459a1-0fad-48a8-88e9-6a145502df04",
   "metadata": {},
   "source": [
    "响应返回的是 LLM 的输出；在这个例子中，由于我们要求模型对一个已经很简短的句子进行总结，它返回了以下消息：\n",
    "\n",
    "`['Michael Smith made a mistake by using PyTorch Trainer instead of HF Trainer.']`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "536aa495-43ab-4570-8142-b040b85f4cf9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'user_input': 'Michael Smith ([REDACTED EMAIL ADDRESS], (+1) [REDACTED PHONE NUMBER]) committed a mistake when he used PyTorch Trainer instead of HF Trainer.\\n\\nSummarize the above text in 1-2 sentences.', 'user_email': None, 'cohere_response': {'data': ['Michael Smith made a mistake by using PyTorch Trainer instead of HF Trainer.'], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 48, 'output_tokens': 14}}}, 'cohere_model': 'command-r-plus', 'temperature': 0.3, 'extras': '{}', 'created_at': datetime.datetime(2024, 6, 10, 2, 16, 7, 666438), 'cohere_endpoint': 'generate'}\n"
     ]
    }
   ],
   "source": [
    "print(db_record)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04dde01f-60f4-4917-a881-3f57baf44c9f",
   "metadata": {},
   "source": [
    "第二个返回项是数据库记录。该代码库是为使用 Postgres 后端而设计的；实际上，代码库自带一个使用 Docker 构建的完整前端。Postgres 数据库用于存储网关的聊天历史记录。然而，它也非常有用，因为它展示了每个请求中实际发送的数据。如我们所见，提示语经过了清洗，实际发送的内容如下：\n",
    "\n",
    "`Michael Smith ([REDACTED EMAIL ADDRESS], (+1) [REDACTED PHONE NUMBER]) committed a mistake when he used PyTorch Trainer instead of HF Trainer.\\n\\nSummarize the above text in 1-2 sentences.`\n",
    "\n",
    "但等等，我听到你在想。Michael Smith 不就是 PII 吗？确实是。但是这个代码库实际上并没有实现姓名清洗功能。接下来，我们将探讨在提示语中应用了哪些清洗器：\n",
    "\n",
    "> [!提示]  \n",
    "> Cohere 的 `generate` 端点实际上已经被弃用，因此，如果有人能为 Cohere API 的新的 Chat 端点创建并提交集成，这将是一个非常棒的开源贡献。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10b182a1-c8a3-4a18-90ae-b8256efea50e",
   "metadata": {},
   "source": [
    "## 清洗器！\n",
    "\n",
    "根据他们的代码库，以下是实现的清洗器：\n",
    "\n",
    "```python\n",
    "ALL_SCRUBBERS = [\n",
    "    scrub_phone_numbers,\n",
    "    scrub_credit_card_numbers,\n",
    "    scrub_email_addresses,\n",
    "    scrub_postal_codes,\n",
    "    scrub_sin_numbers,\n",
    "]\n",
    "```\n",
    "\n",
    "网关会依次应用每个清洗器。\n",
    "\n",
    "这虽然有些 “hacky”（不够优雅），但如果你确实需要实现另一个清洗器，可以通过修改包装器方法来实现。以下是一个演示：\n",
    "\n",
    "> [!提示]  \n",
    "> 作者提到，SIN（社会保险号）清洗器特别容易误清洗数据，因此它会被放在最后，以确保其他与数字相关的 PII 先被清洗。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "e5ca6d12-0663-4311-9b6f-e3596fc2a36e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_custom_scrubber(text: str) -> str:\n",
    "    \"\"\"\n",
    "    Scrub Michael Smith in text\n",
    "\n",
    "    :param text: Input text to scrub\n",
    "    :type text: str\n",
    "    :return: Input text with any mentions of Michael Smith scrubbed\n",
    "    :rtype: str\n",
    "    \"\"\"\n",
    "    return re.sub(\n",
    "        r\"Michael Smith\",\n",
    "\n",
    "        \n",
    "        \"[REDACTED PERSON]\",\n",
    "        text,\n",
    "        re.IGNORECASE\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "1c31aed5-f934-4498-8a91-e509eb041a91",
   "metadata": {},
   "outputs": [],
   "source": [
    "original_method = wrapper.send_cohere_request\n",
    "\n",
    "def modified_method(self, **kwargs):\n",
    "    self._validate_cohere_endpoint(kwargs.get('endpoint', None)) # Unfortunate double validate cohere endpoint call\n",
    "    prompt = kwargs.get('prompt', None)\n",
    "    text = my_custom_scrubber(prompt)\n",
    "    kwargs['prompt'] = text\n",
    "    return original_method(**kwargs)\n",
    "\n",
    "# Assign the new method to the instance\n",
    "wrapper.send_cohere_request = types.MethodType(modified_method, wrapper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "id": "3434529d-517d-48e6-a891-35108366fb90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data': ['[REDACTED PERSON] made an error by using PyTorch Trainer instead of HF Trainer. They can be contacted at [RED'], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 52, 'output_tokens': 25}}}\n"
     ]
    }
   ],
   "source": [
    "response, db_record = wrapper.send_cohere_request(\n",
    "    endpoint=\"generate\",\n",
    "    model=\"command-r-plus\",\n",
    "    max_tokens=25,\n",
    "    prompt=f\"{example}\\n\\nSummarize the above text in 1-2 sentences.\",\n",
    "    temperature=0.3,\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "b4fcd05a-6648-4a15-9195-ab865b2d5f05",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'user_input': '[REDACTED PERSON] ([REDACTED EMAIL ADDRESS], (+1) [REDACTED PHONE NUMBER]) committed a mistake when he used PyTorch Trainer instead of HF Trainer.\\n\\nSummarize the above text in 1-2 sentences.', 'user_email': None, 'cohere_response': {'data': ['[REDACTED PERSON] made an error by using PyTorch Trainer instead of HF Trainer. They can be contacted at [RED'], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 52, 'output_tokens': 25}}}, 'cohere_model': 'command-r-plus', 'temperature': 0.3, 'extras': '{}', 'created_at': datetime.datetime(2024, 6, 10, 2, 59, 58, 733195), 'cohere_endpoint': 'generate'}\n"
     ]
    }
   ],
   "source": [
    "print(db_record)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c8da471-285c-4c6f-9d37-d27eb5479a6b",
   "metadata": {},
   "source": [
    "如果你确实需要这样做，请务必记住，清洗器是按顺序应用的，因此，如果你的自定义清洗器与任何默认清洗器发生冲突，可能会导致一些意外行为。\n",
    "\n",
    "例如，针对姓名的清洗，实际上有[其他清洗库](https://github.com/kylemclaren/scrub)可以探索，这些库采用更复杂的算法来清洗 PII。这个代码库涵盖了更多的PII类型，例如 [IP 地址、主机名等](https://github.com/kylemclaren/scrub/blob/master/scrubadubdub/scrub.py)。然而，如果你仅仅需要删除特定的匹配项，你仍然可以使用上述代码进行处理。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf338c5d-cd2c-4f9a-9ea2-ea18b352c60c",
   "metadata": {},
   "source": [
    "## 数据集\n",
    "\n",
    "让我们在一个完整的数据集上演示这个包装器的实际应用。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "47166d9f-d4e1-4895-b3b9-9a5fbb399e9c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "409122f13cd748c78f9b1be3cbfac4f8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/12.8k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a61bbc66ab649e39af020347df2331a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/73.8M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2f33d9cb809c4aada68fc35a7a2b68a1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/116M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b358959872654504b91421352adfd3b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/97.8M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "46a96422ad954882b1658f211521bf16",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/93.1M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c6771a1fdfbd40a488bd4fab81cc88b0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pii_ds = load_dataset(\"ai4privacy/pii-masking-200k\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "7df53225-9880-4f04-8c8e-bed5de6aa52c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"I need the latest update on assessment results. Please send the files to Valentine4@gmail.com. For your extra time, we'll offer you Kip 100,000 but please provide your лв account details.\""
      ]
     },
     "execution_count": 141,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pii_ds['train'][36]['source_text']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "d2f1c854-7fe3-4db3-8b88-635dc5894b38",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data': [\"The person is requesting an update on assessment results and is offering Kip 100,000 in exchange for the information and the recipient's account details.\"], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 64, 'output_tokens': 33}}}\n"
     ]
    }
   ],
   "source": [
    "example = pii_ds['train'][36]['source_text']\n",
    "\n",
    "response, db_record = wrapper.send_cohere_request(\n",
    "    endpoint=\"generate\",\n",
    "    model=\"command-r-plus\",\n",
    "    max_tokens=50,\n",
    "    prompt=f\"{example}\\n\\nSummarize the above text in 1-2 sentences.\",\n",
    "    temperature=0.3,\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "2c8189d8-9878-430b-b053-0d340acbe008",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'user_input': \"I need the latest update on assessment results. Please send the files to V[REDACTED EMAIL ADDRESS]. For your extra time, we'll offer you Kip 100,000 but please provide your лв account details.\\n\\nSummarize the above text in 1-2 sentences.\", 'user_email': None, 'cohere_response': {'data': [\"The person is requesting an update on assessment results and is offering Kip 100,000 in exchange for the information and the recipient's account details.\"], 'return_likelihoods': None, 'meta': {'api_version': {'version': '1'}, 'billed_units': {'input_tokens': 64, 'output_tokens': 33}}}, 'cohere_model': 'command-r-plus', 'temperature': 0.3, 'extras': '{}', 'created_at': datetime.datetime(2024, 6, 10, 3, 10, 51, 416091), 'cohere_endpoint': 'generate'}\n"
     ]
    }
   ],
   "source": [
    "print(db_record)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aab8583-daf4-4374-89b9-b596715425b8",
   "metadata": {},
   "source": [
    "## 常规输出\n",
    "\n",
    "如果我们直接将文本发送到端点而不进行任何清洗，摘要结果如下所示："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "id": "d381c8ea-4105-4def-878e-634100ea7d78",
   "metadata": {},
   "outputs": [],
   "source": [
    " co = cohere.Client(\n",
    "    api_key=os.environ['COHERE_API_KEY']\n",
    ")\n",
    "\n",
    "response_vanilla = co.generate(\n",
    "    prompt=f\"{example}\\n\\nSummarize the above text in 1-2 sentences.\",\n",
    "    model=\"command-r-plus\",\n",
    "    max_tokens=50,\n",
    "    temperature=0.3\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "15fad584-75f3-4955-bd31-06e9adff9dd1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>prompt</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>I need the latest update on assessment results. Please send the files to Valentine4@gmail.com. For your extra time, we'll offer you Kip 100,000 but please provide your лв account details.\n",
       "\n",
       "Summarize the above text in 1-2 sentences.</td>\n",
       "      <td>The text is a request for an update on assessment results to be sent to Valentine4@gmail.com, with an offer of Kip 100,000 in exchange for the information and account details.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "Generations([cohere.Generation {\n",
       "             \tid: f3b759b4-2a58-467c-af9d-288e769a5a44\n",
       "             \tprompt: I need the latest update on assessment results. Please send the files to Valentine4@gmail.com. For your extra time, we'll offer you Kip 100,000 but please provide your лв account details.\n",
       "             \n",
       "             Summarize the above text in 1-2 sentences.\n",
       "             \ttext: The text is a request for an update on assessment results to be sent to Valentine4@gmail.com, with an offer of Kip 100,000 in exchange for the information and account details.\n",
       "             \tlikelihood: None\n",
       "             \tfinish_reason: COMPLETE\n",
       "             \ttoken_likelihoods: None\n",
       "             }])"
      ]
     },
     "execution_count": 148,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response_vanilla"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7103c697-6479-4907-a646-ed5ac2c7dd09",
   "metadata": {},
   "source": [
    "总结一下，在 Notebook 中，我们演示了如何使用 Wealthsimple 开源的 PII 检测示例网关，并在此基础上添加了自定义清洗器。如果你真的需要可靠的 PII 检测，确保运行自己的测试，验证你所采用的清洗算法是否真正覆盖了你的应用场景。最重要的是，尽可能在你自己托管的基础设施上部署开源模型，这将始终是构建 LLM 应用时最安全、最可靠的选择 :)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nlp",
   "language": "python",
   "name": "nlp"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
